#Author: Mike Schladt - 2014
#Controls system processes of the MIP

import subprocess 
import time
import os
import re
import signal
from termcolor import colored

#global dirs
home_dir = os.environ['HOME']
cuckoo_dir = '{0}/cuckoo/'.format(home_dir)
utils_dir = '{0}/cuckoo/utils/'.format(home_dir)
mip_dir = '{0}/malware-intelligence-project/'.format(home_dir)
soft_timeout = 300
hard_timeout = 500 

def check_health():
    #check for running processes
    #return False if not healthy
     
    cuckoo_on = False
    api_on = False
    controller_on = False

    s = subprocess.Popen(["ps", "axw"],stderr=subprocess.STDOUT,stdout = subprocess.PIPE)
    for x in s.stdout:
        #search for required running processes by NAME!!!
        #required because gnome-terminal doesn't terminate as expected
        if re.search('python cuckoo.py',x) and not re.search('-e',x) :
            cuckoo_on = True
        if re.search('python api.py',x) and not re.search('-e',x) :
            api_on = True
        if re.search('controller.py',x) :
            controller_on = True
        if cuckoo_on and api_on and controller_on :
            return True
    return False

def reboot():
    #kills all processes; cleans up; reboot
    try:
        p1.kill()
    except: 
        pass
    try:
        p2.kill()
    except:
        pass
    try:
        p3.kill()
    except: 
        pass

    subprocess.call(['gnome-terminal', '--tab', '--working-directory={0}'.format(utils_dir), '-e', 'bash ./clean.sh'])
    subprocess.call(['pkill', 'VBoxHeadless'])
    time.sleep(10)

##############################################################

#run the init script to turn on required services
subprocess.call(['/bin/sh', '{0}init.sh'.format(mip_dir)])

#start and maintain main system processes 
while 1:
    master_counter = 0
    p1 = subprocess.Popen(['gnome-terminal', '--tab', '--working-directory={0}'.format(cuckoo_dir), '-e', 'python cuckoo.py']) 
    p2 = subprocess.Popen(['gnome-terminal', '--tab', '--working-directory={0}'.format(utils_dir), '-e', 'python api.py']) 
    time.sleep(6)    
    p3 = subprocess.Popen(['python', '{0}controller.py'.format(mip_dir)])

    #main loop to check health and after given time
    while (check_health() and master_counter< hard_timeout):
        print "####################### OVERLORD LOOP {0} #######################".format(master_counter)
        if (master_counter == soft_timeout):
            print(colored("Hitting soft-timout: sending signal...","red"))
            p3.send_signal(signal.SIGINT)         
        master_counter += 1    
        time.sleep(6)

    reboot()


